from fusionmatrix.analyzers.classifier.template import MULTI_CLASSIFICATION_TEMPLATE, CLASSIFICATION_TEMPLATE
from fusionmatrix.base.task import BaseTask, TaskResult
from fusionmatrix.base.model import ChatModel
from loguru import logger
from collections import OrderedDict
from pydantic import BaseModel, Field
from typing import List, Dict, Tuple


class Option(BaseModel):
    name: str
    explanation: str


class ClassificationResult(BaseModel):
    succeeded: bool = True
    classes: List[str] = Field(default_factory=list)
    message: str = ''

    @classmethod
    def failure(cls, reason: str) -> 'ClassificationResult':
        return ClassificationResult(succeeded=False, message=reason)


class Classifier(BaseModel):
    name: str = ''
    description: str = ''
    options: List[Option] = Field(default_factory=list)
    related_columns: List[str] = Field(default_factory=list)
    multi: bool = False

    def hasOption(self, option_name: str) -> bool:
        for option in self.options:
            if option.name == option_name:
                return True

        return False

    def getClassificationPrompt(self, content: str) -> str:
        options = '\n'.join(f'- {option.name}: {option.explanation}' for option in self.options)
        if self.multi:
            return MULTI_CLASSIFICATION_TEMPLATE.format(
                content=content,
                options=options
            )
        else:
            return CLASSIFICATION_TEMPLATE.format(
                content=content,
                options=options
            )

    def parseAnswer(self, answer: str) -> ClassificationResult:
        lines = list(filter(lambda line: line.strip() != '', answer.splitlines()))
        if len(lines) != 2:
            return ClassificationResult.failure("输出行数不符")
        elif ':' not in lines[1]:
            return ClassificationResult.failure("第二行未按照给定格式输出")

        result = ClassificationResult()
        result.message = lines[0]
        alternave_classes = lines[1].split(':')[1].split(',')
        if len(alternave_classes) != 1 and self.multi is False:
            return ClassificationResult.failure(f"对于一个单分类任务输出了多个分类")

        for alternave_class in alternave_classes:
            alternave_class = alternave_class.strip()
            if not self.hasOption(alternave_class):
                return ClassificationResult.failure(f"输出的分类结果中，'{alternave_class}'不属于给定的类别")

            result.classes.append(alternave_class)

        return result


class Classifiers(BaseModel):
    classifiers: List[Classifier]


class ClassificationTask(BaseTask):
    classifier: Classifier
    model: ChatModel
    target: Tuple[str, int, int]
    row: OrderedDict[str, str] = Field(default_factory=OrderedDict)

    def run(self) -> TaskResult:
        content = ""
        for col_name in self.classifier.related_columns:
            col_content = self.row[col_name].strip()
            if col_content != "":
                content += f"- [{col_name}] 列的内容为：{col_content}\n"

        if content == "":
            # 没有有效的分类信息，不需要进行分类
            return TaskResult(task_id=self.task_id, task_status=True, result='')

        initial_question = self.classifier.getClassificationPrompt(content)
        instance = self.model.getInstance()
        answer = instance.ask(initial_question)
        ntry = 0

        result = self.classifier.parseAnswer(answer)
        while result.succeeded is False and ntry <= 5:
            ntry += 1
            logger.debug(answer)
            logger.warning(f"task {self.task_id} answer invalid format: {result.message}")
            answer = instance.ask(f"你的输出不符合要求：{result.message}，请根据给定的格式要求重新输出")
            result = self.classifier.parseAnswer(answer)

        return TaskResult(task_id=self.task_id, task_status=True, result=','.join(result.classes))

    def notice(self, result: TaskResult) -> None:
        from fusionmatrix.windows.main import MainWindow
        if result.task_status:
            table = MainWindow.getTableByName(self.target[0])
            table.setCellText(self.target[1], self.target[2], result.result)
